//
//  keypoint_detector.cpp
//  yuzon
//
//  Created by Alvin on 2023/3/29.
//

#include "keypoint_detector.hpp"
#include "transforms.hpp"
#include "utils.hpp"

using namespace cv;
using namespace std;
using namespace paddle::lite_api;

KeypointDetector::KeypointDetector(const string &model_keypoint_path) :
input_width_(kKeypointInputWidth),
input_height_(kKeypointInputHeight),
dark_(kDark) {
  MobileConfig config;
  config.set_model_from_file(model_keypoint_path);
  config.set_power_mode(LITE_POWER_NO_BIND);
  config.set_threads(4);
  predictor_ = CreatePaddlePredictor<MobileConfig>(config);
}

void KeypointDetector::Predict(const Mat &image,
                               const HumanResults &human_results,
                               KeypointResults &keypoint_results) {
  if (human_results.empty())
    return;
  
  vector<Mat> crops;
  vector<Point2f> centers;
  vector<Size2f> scales;
  
  CropImage(image, human_results, crops, centers, scales);
  Preprocess(crops);
  predictor_->Run();
  Postprocess(image, centers, scales, keypoint_results);
}

void KeypointDetector::CropImage(const Mat &image,
                                 const HumanResults &results,
                                 vector<Mat> &crops,
                                 vector<Point2f> &centers,
                                 vector<Size2f> &scales) {
  for (const HumanResult &result : results) {
    const Rect &rect = result.rect;
    
    int x = rect.x;
    int y = rect.y;
    int half_w = round(rect.width / 2);
    int half_h = round(rect.height / 2);
    int center_x = x + half_w;
    int center_y = y + half_h;
    
    // adjust w or h to keep image ratio, expand the shorter edge
    if (half_h * input_width_ > half_w * input_height_)
      half_w = round(half_h * input_width_ / input_height_);
    else
      half_h = round(half_w * input_height_ / input_width_);
    
    int roi_x1 = center_x - half_w;
    int roi_y1 = center_y - half_h;
    int roi_x2 = center_x + half_w;
    int roi_y2 = center_y + half_h;
    
    int crop_x1 = max(roi_x1, 0);
    int crop_y1 = max(roi_y1, 0);
    int crop_x2 = min(roi_x2, image.cols);
    int crop_y2 = min(roi_y2, image.rows);
    
    Mat crop = image(Range(crop_y1, crop_y2), Range(crop_x1, crop_x2));
    Mat roi;
    copyMakeBorder(crop, roi,
                   max(0, -roi_y1), max(0, roi_y2 - image.rows),
                   max(0, -roi_x1), max(0, roi_x2 - image.cols),
                   BORDER_CONSTANT, Scalar(0, 0, 0));
    Point2f center(center_x, center_y);
    Size2f scale(half_w * 2, half_h * 2);
    
    crops.emplace_back(roi);
    centers.emplace_back(center);
    scales.emplace_back(scale);
  }
}

void KeypointDetector::Preprocess(const vector<Mat> &crops) {
  int64_t num_crops = crops.size();
  
  auto input_tensor = predictor_->GetInput(0);
  input_tensor->Resize({num_crops, 3, input_height_, input_width_});
  auto input_data = input_tensor->mutable_data<float>();
  
  for (int i = 0; i < num_crops; ++i) {
    Mat input_image;
    resize(crops[i], input_image, Size(input_width_, input_height_));
    input_image.convertTo(input_image, CV_32FC3);
    Permute(input_image, input_data + i * (3 * input_height_ * input_width_));
  }
}

void KeypointDetector::Postprocess(const Mat &image,
                                   const vector<Point2f> &centers,
                                   const vector<Size2f> &scales,
                                   KeypointResults &results) {
  auto output_tensor = predictor_->GetOutput(0);
  auto output_data = output_tensor->data<float>();
  vector<int64_t> output_shape = output_tensor->shape();
  
  int64_t outputs_size = ShapeProduct(output_shape);
  vector<float> outputs(outputs_size);
  copy_n(output_data, outputs_size, outputs.data());
  
  auto index_tensor = predictor_->GetOutput(1);
  auto index_data = index_tensor->data<int64_t>();
  vector<int64_t> index_shape = index_tensor->shape();
  
  int64_t indexes_size = ShapeProduct(index_shape);
  vector<int64_t> indexes(indexes_size);
  copy_n(index_data, indexes_size, indexes.data());
  
  for (int i = 0; i < output_shape[0]; ++i) {
    int64_t output_size = output_shape[1] * output_shape[2] * output_shape[3];
    vector<float> output(outputs.begin() + i * output_size,
                         outputs.begin() + (i + 1) * output_size);
    int64_t index_size = output_shape[1];
    vector<int64_t> index(indexes.begin() + i * index_size,
                          indexes.begin() + (i + 1) * index_size);
    vector<float> preds(output_shape[1] * 3);
    
    get_final_preds(output, output_shape, index,
                    centers[i], scales[i], preds, dark_);
    
    KeypointResult result;
    result.num_joints = int(output_shape[1]);
    
    for (int i = 0; i < output_shape[1]; ++i) {
      result.keypoints.emplace_back(preds[i * 3]);
      result.keypoints.emplace_back(preds[i * 3 + 1]);
      result.keypoints.emplace_back(preds[i * 3 + 2]);
    }
    
    results.emplace_back(result);
  }
}
